#define WIN32_LEAN_AND_MEAN
#define _NO_CVCONST_H
#include <Windows.h>
#include <DbgHelp.h>
#include <stdio.h>
#include <stdlib.h>


#define MAX_CHILD_COUNT 128


// Taken from cvconst.h from the DIA SDK
enum BasicType
{
    btNoType = 0,
    btVoid = 1,
    btChar = 2,
    btWChar = 3,
    btInt = 6,
    btUInt = 7,
    btFloat = 8,
    btBCD = 9,
    btBool = 10,
    btLong = 13,
    btULong = 14,
    btCurrency = 25,
    btDate = 26,
    btVariant = 27,
    btComplex = 28,
    btBit = 29,
    btBSTR = 30,
    btHresult = 31,

    btMax
};
enum UdtKind
{
    UdtStruct,
    UdtClass,
    UdtUnion
};


typedef struct
{
    const char* name;
    DWORD64 base;
    DWORD size;
} Module;

HANDLE process = NULL;
Module* modules = NULL;
size_t moduleCount = 0;
int displayFunctions = 0;
const char* basicTypeNames[btMax];


void DisplayHelp()
{
    printf("addr2line.exe - Resolves addresses into file names, line numbers and optionally function names.\n");
    printf("addr2line.exe [-f]\n");
    printf("              [-e filename]\n");
    printf("              [-b base]\n");
    printf("              [-s size]\n");
    printf("\n");
    printf("[-f]\n");
    printf("   Output function names as well as file and line number information.\n");
    printf("[-e filename]\n");
    printf("   Specify a module from which symbols should be extracted. Multiple instances are accepted.\n");
    printf("[-b base]\n");
    printf("   Specify the base address for the preceeding module. Multiple instances are accepted.\n");
    printf("[-s size]\n");
    printf("   Specify the size of the preceeding module. Multiple instances are accepted.\n");
    printf("\n");
    printf("Example usage: addr2line.exe -f -e foo.exe -b 0x00bc0000 -s 0x001eb000 -e ntdll.dll -b 0x76ec0000 -s 0x00127000\n");
}


int ParseArguments(int argc, char* argv[])
{
    int i;
    if (argc < 2)
        return 1;

    for (i=1; i<argc; ++i)
    {
        const char* arg = argv[i];
        if (!strcmp(arg, "-f"))
        {
            displayFunctions = 1;
        }
        else if (!strcmp(arg, "-e"))
        {
            if (++i == argc) return 2;
            modules[moduleCount++].name = argv[i];
        }
        else if (!strcmp(arg, "-b"))
        {
            if (++i == argc) return 2;
            if (moduleCount == 0) return 3;
            if (!sscanf_s(argv[i], "0x%xl", &modules[moduleCount-1].base) && !sscanf_s(argv[i], "0X%xl", &modules[moduleCount-1].base))
                sscanf_s("%ul", argv[i], &modules[moduleCount-1].base);
        }
        else if (!strcmp(arg, "-s"))
        {
            if (++i == argc) return 2;
            if (moduleCount == 0) return 3;
            if (!sscanf_s(argv[i], "0x%x", &modules[moduleCount-1].size) && !sscanf_s(argv[i], "0X%x", &modules[moduleCount-1].size))
                sscanf_s("%u", argv[i], &modules[moduleCount-1].size);
        }
        else
        {
            return 4;
        }
    }

    return 0;
}


void CleanupAndExit(int err, size_t modulesToUnloadCount)
{
    size_t i;
    for (i=0; i<modulesToUnloadCount; ++i)
        SymUnloadModule64(process, modules[i].base);
    SymCleanup(process);
    free(modules);
    exit(err);
}


int main(int argc, char* argv[])
{
    SYMBOL_INFO* symbol;
    TI_FINDCHILDREN_PARAMS* children;
    IMAGEHLP_LINE64 line;
    size_t i;
    int result;

    process = GetCurrentProcess();

    // Allocate enough space to hold all possible modules being passed in.
    // This over allocates a bit, but it's simpler this way.
    modules = (Module*)malloc(argc*sizeof(Module));
    memset(modules, 0, argc*sizeof(Module));

    result = ParseArguments(argc, argv);
    if (result != 0)
    {
        DisplayHelp();
        CleanupAndExit(result, moduleCount);
    }

    // Initialize the symbol handler
    if (!SymInitialize(process, NULL, FALSE))
        CleanupAndExit(5, moduleCount);
    SymSetOptions(SYMOPT_LOAD_LINES | SYMOPT_FAIL_CRITICAL_ERRORS | SYMOPT_UNDNAME);

    // Load all the modules passed in
    for (i=0; i<moduleCount; ++i)
    {
        DWORD64 address = SymLoadModuleEx(process, NULL, modules[i].name, NULL, modules[i].base, modules[i].size, NULL, 0);
        if (address == 0)
            CleanupAndExit((int)GetLastError(), i);
    }

    // Initialize the structures used for extracting symbol information
    symbol = (SYMBOL_INFO*)malloc(sizeof(SYMBOL_INFO) + MAX_SYM_NAME*sizeof(TCHAR));
    memset(symbol, 0, sizeof(SYMBOL_INFO) + MAX_SYM_NAME*sizeof(TCHAR));
    symbol->SizeOfStruct = sizeof(SYMBOL_INFO);
    symbol->MaxNameLen = MAX_SYM_NAME;
    memset(&line, 0, sizeof(IMAGEHLP_LINE64));
    line.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
    children = (TI_FINDCHILDREN_PARAMS*)malloc(sizeof(TI_FINDCHILDREN_PARAMS) + MAX_CHILD_COUNT*sizeof(ULONG));
    memset(children, 0, sizeof(TI_FINDCHILDREN_PARAMS) + MAX_CHILD_COUNT*sizeof(ULONG));

    // Start interactive symbol lookup
    for (;;)
    {
        DWORD64 address;
        DWORD offset;
        char arg[512];

        // Read and parse an address in either hexadecimal, or decimal format
        if (scanf_s("%s", arg, 512) != 1)
            break;
        if (!sscanf_s(arg, "0x%I64x", &address) && !sscanf_s(arg, "0X%I64x", &address))
            sscanf_s("%I64u", arg, &address);

        // Lookup funcion name if required
        if (displayFunctions)
        {
            DWORD64 d;
            if (SymFromAddr(process, address, &d, symbol))
            {
                printf("%s\n", symbol->Name);
            }
            else
                printf("??\n");
        }

        // Lookup file and line number
        if (SymGetLineFromAddr64(process, address, &offset, &line))
            printf("%s:%d\n", line.FileName, line.LineNumber);
        else
            printf("??:0\n");

        // Flush the output, to play nicely with any other processes who might be redirecting our output
        fflush(stdout);
    }

    // Clean up
    free(symbol);
    CleanupAndExit(0, moduleCount);

    return 0;
}
